"""
trimesh pyradiomics package is needed
"""
import os
from django.core.wsgi import get_wsgi_application  # noqa
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'biomind.settings')  # noqa
os.environ["DJANGO_ALLOW_ASYNC_UNSAFE"] = "true"  # noqa
application = get_wsgi_application()  # noqa
from bioutils.dicom.constants import ImageFolderType
from bioutils.transformer import to_npy, DataTransformer
from bioutils.files import FileTool
from rest_framework import serializers
from predicts.models import Study, Tasks, Predicts, Segments, Masks, Series
from com_models.types import TaskStatus, TaskDeleteStatus
from conf.globals import CACHE_MODE, CACHE_PATH, Predictors
from datetime import datetime
import numpy as np
import SimpleITK as sitk
from radiomics import featureextractor
from biomind.settings import BASE_DIR
import pandas as pd
from common.log import logger
from django.db.models import Q
from common.decorators.stopwatch import stopwatch

PARAMS = (BASE_DIR / "scripts/feature_extractor/param.yaml").as_posix()
# PARAMS = FileTool.read_yaml(params_path)


class MasksSerializer(serializers.ModelSerializer):
    sop_instance_uid = serializers.CharField()
    mask = serializers.SerializerMethodField()

    def get_mask(self, instance):
        # 需要对绝对路径进行切割返回
        return FileTool.read_png(instance.mask_path)

    class Meta:
        model = Masks
        fields = ['sop_instance_uid', 'mask_path', 'mask']


class StudySerializer(serializers.ModelSerializer):
    class Meta:
        model = Study
        fields = ['patient_id', 'patient_name', 'accession_number', 'study_instance_uid']


class SeriesSerializer(serializers.ModelSerializer):
    class Meta:
        model = Series
        fields = ['series_number', 'series_instance_uid']


def get_conteverted_features(features):
    feature_dict = {}
    for k, v in features.items():
        if not k.startswith("diagnostics_"):
            feature_dict[k] = v
    return feature_dict


segment_map = {
    "inflammatory_lesions": '炎性',
    "fiber": '索条性',
    "infectious": "series"
}


@stopwatch('')
def extract_features(start_datetime=None, end_datetime=None, csv_name="res.csv"):
    start_datetime = start_datetime or datetime(1900, 1, 1, 0, 0, 0)
    end_datetime = end_datetime or datetime.now()
    study_infos = Study.objects.filter(study_datetime__range=(start_datetime, end_datetime))
    extractor = featureextractor.RadiomicsFeatureExtractor(PARAMS)
    result = []
    for study_info in study_infos:
        study_uid = study_info.study_instance_uid
        predict_objs = Predicts.objects.filter(
            study_uid=study_uid,
            task_id__status=TaskStatus.complete,
            task_id__is_active=TaskDeleteStatus.effective,
            predictor=Predictors.lungct_v2_predictor)
        study_dict = dict(StudySerializer(study_info).data)
        for predict_obj in predict_objs:
            for volume_info in predict_obj.prediction.get('classifier_series', {}).values():
                series_uid = volume_info['series_uid']
                vol_path = volume_info['vol_path']
                vol_sort_path = volume_info['vol_sort_path']
                segment_objs = Segments.objects.filter(
                    series_uid=series_uid, dkey='infectious')
                if not segment_objs.exists():
                    continue
                image_sitk = DataTransformer.MhdToSitk(vol_path)
                image_arrs = to_npy(vol_path, ImageFolderType.Mhd)
                image_sorts = FileTool.read_json(vol_sort_path)
                image_arrs = DataTransformer.SitkToNpy(image_sitk)
                series_dict = dict(SeriesSerializer(Series.objects.get(series_instance_uid=series_uid)).data)
                logger.info(
                    f"patient_id: {study_info.patient_id}, predictor: {predict_obj.predictor}, predict_id: {predict_obj.id}")

                series_mask_arrs = np.zeros_like(image_arrs)
                for segment_obj in segment_objs:
                    mask_sers = MasksSerializer(instance=Masks.objects.filter(segment_id=segment_obj), many=True)
                    name = segment_obj.extensions.get('name')
                    logger.info(
                        f"calculate single segment for series_num: {series_dict['series_number']}, name: {name}")
                    mask_arrs = np.zeros_like(image_arrs)
                    mask_objs = mask_sers.data
                    for mask_obj in mask_objs:
                        index = image_sorts.index(mask_obj['sop_instance_uid'])
                        mask_arrs[index] = mask_obj['mask']
                        series_mask_arrs[index] = mask_obj['mask']
                    mask_sitk = sitk.GetImageFromArray(mask_arrs)
                    mask_sitk.CopyInformation(image_sitk)
                    features = extractor.execute(image_sitk, mask_sitk)
                    feature_dict = get_conteverted_features(features)
                    result.append({**study_dict,
                                   **series_dict,
                                   "segment": segment_map.get(name),
                                   **feature_dict})

                logger.info(
                    f"calculate all series segment for series_num: {series_dict['series_number']}, dkey: {segment_obj.dkey}")
                series_mask_sitk = sitk.GetImageFromArray(series_mask_arrs)
                series_mask_sitk.CopyInformation(image_sitk)
                features = extractor.execute(image_sitk, series_mask_sitk)
                feature_dict = get_conteverted_features(features)
                result.append({**study_dict,
                               **series_dict,
                               "segment": segment_map.get(segment_obj.dkey),
                               **feature_dict})

    df = pd.DataFrame(result)
    df.to_csv(csv_name, index=False, encoding="utf-8")
    print('-' * 10)


if __name__ == '__main__':
    extract_features(datetime(2022, 1, 1, 0, 0, 0), datetime(2022, 1, 18, 0, 0, 0))
    # extract_features()
